package test;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.CipherOutputStream;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import util.FileUtil;

public class RSA_AES_TEST {

	private static KeyPairGenerator keyGen;
	private static KeyPair generatedKey;

	private static PrivateKey localPrivateKey;
	private static PublicKey distantPublicKey;

	private static SecretKey secretKeyAES;
	private static SecretKey decryptedKey;

	public static void main(String[] args) {

		generatedKey = generateKeys();

		localPrivateKey = generatedKey.getPrivate();
		distantPublicKey = generatedKey.getPublic();

		secretKeyAES = getAESKey();
		
		byte[] secretKeyAESBytes = EncryptSecretKey();

		//Load file
		String path =  "C:/Users/Alex/Pictures/Wallpaper/fantasy-street-abstract-nature-night-city-landscape-dream-wallpapers-wallpaper-wallwuzz-hd-wallpaper-10282.jpg";
		byte[] fileData = FileUtil.fileToBytes(path);

		System.out.println("Starting AES encryption");

		byte[] encryptedData = encrypt(fileData, secretKeyAES);

		System.out.println("Finshed encrypting with AES");

		System.out.println("Starting AES decryption");
		
		byte[] copiedEncryptedData = new byte[encryptedData.length];
		
		for(int i = 0; i < copiedEncryptedData.length; i++) {
			copiedEncryptedData[i] = encryptedData[i];
		}
		
		decryptedKey = decryptAESKey(secretKeyAESBytes);

		byte[] decryptedData = decrypt(copiedEncryptedData, decryptedKey);

		System.out.println("Finished AES decryption");

		FileUtil.bytesToFile("C:/Users/Alex/LOG735-Client/fantasy-city.jpg", decryptedData);

	}

	public static byte[] encrypt(byte[] data, SecretKey secretKey) {
		Cipher encryptCipher = null;
		byte[] encryptedData = null;
		try {
			encryptCipher = Cipher.getInstance("AES");
			encryptCipher.init(Cipher.ENCRYPT_MODE, secretKey);
			encryptedData = encryptCipher.doFinal(data);
		} catch (InvalidKeyException | NoSuchAlgorithmException | NoSuchPaddingException | IllegalBlockSizeException | BadPaddingException e) {
			e.printStackTrace();
		}
		return encryptedData;
	}

	public static byte[] decrypt(byte[] encryptedData, SecretKey secretKey) {
		Cipher decryptCipher = null;
		byte[] decryptedData = null;
		try {
			decryptCipher = Cipher.getInstance("AES");
			decryptCipher.init(Cipher.DECRYPT_MODE, secretKey);
			decryptedData = decryptCipher.doFinal(encryptedData);
		} catch (InvalidKeyException | NoSuchAlgorithmException | NoSuchPaddingException | IllegalBlockSizeException | BadPaddingException e) {
			e.printStackTrace();
		}
		return decryptedData;
	}


	private static byte[] EncryptSecretKey()
	{
		Cipher cipher = null;
		byte[] key = null;

		try
		{
			// initialize the cipher with the user's public key
			cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
			cipher.init(Cipher.ENCRYPT_MODE, distantPublicKey);
			key = cipher.doFinal(secretKeyAES.getEncoded());
		}
		catch(Exception e )
		{
			System.out.println ( "exception encoding key: " + e.getMessage() );
			e.printStackTrace();
		}
		return key;
	}

	private static SecretKey decryptAESKey(byte[] data)
	{
		SecretKey key = null;
		PrivateKey privKey = null;
		Cipher cipher = null;

		try
		{
			// this is OUR private key
			privKey = localPrivateKey;

			// initialize the cipher...
			cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
			cipher.init(Cipher.DECRYPT_MODE, privKey );

			// generate the aes key!
			key = new SecretKeySpec (cipher.doFinal(data), "AES");
		}
		catch(Exception e)
		{
			System.out.println ( "exception decrypting the aes key: " 
					+ e.getMessage() );
			return null;
		}

		return key;

	}

	private static KeyPair generateKeys() {

		if (keyGen == null){
			try {
				keyGen = KeyPairGenerator.getInstance("RSA");
				SecureRandom random = SecureRandom.getInstance("SHA1PRNG", "SUN");
				keyGen.initialize(1024, random);	
			} catch (NoSuchAlgorithmException | NoSuchProviderException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		}

		return keyGen.generateKeyPair();
	}

	private static SecretKey getAESKey() {
		KeyGenerator keyGen = null;
		try {
			keyGen = KeyGenerator.getInstance("AES");
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		}
		keyGen.init(128);
		SecretKey secretKey = keyGen.generateKey();
		return secretKey;
	}

}
